package com.yeskery.nut.util;

import javax.crypto.Cipher;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Base64;

/**
 * RSA算法工具类
 * @author sunjay
 * 2023/10/9
 */
public class RsaUtils {

    /**
     * 生成密钥对
     * @return 密钥对
     * @throws Exception 发生异常
     */
    public static KeyPair generateKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048);
        return keyPairGenerator.generateKeyPair();
    }

    /**
     * 生成密钥对
     * @param secureRandom 安全的随机数
     * @return 密钥对
     * @throws Exception 发生异常
     */
    public static KeyPair generateKeyPair(SecureRandom secureRandom) throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048, secureRandom);
        return keyPairGenerator.generateKeyPair();
    }

    /**
     * 生成密钥对
     * @param salt 盐
     * @return 密钥对
     * @throws Exception 发生异常
     */
    public static KeyPair generateKeyPair(String salt) throws Exception {
        return generateKeyPair(new SecureRandom(salt.getBytes()));
    }

    /**
     * 使用公钥加密文本
     * @param plainText 文本
     * @param publicKey 公钥
     * @return 加密后的文本
     * @throws Exception 发生异常
     */
    public static byte[] encrypt(String plainText, PublicKey publicKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        return cipher.doFinal(plainText.getBytes());
    }

    /**
     * 使用私钥解密文本
     * @param encryptedBytes 要解密的文本字节数组
     * @param privateKey 私钥
     * @return 解密后的文本
     * @throws Exception 发生异常
     */
    public static byte[] decrypt(byte[] encryptedBytes, PrivateKey privateKey) throws Exception {
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        return cipher.doFinal(encryptedBytes);
    }

    /**
     * 使用私钥对文本进行签名
     * @param plainText 文本
     * @param privateKey 私钥
     * @return 签名后的字节数组
     * @throws Exception 发生异常
     */
    public static byte[] sign(String plainText, PrivateKey privateKey) throws Exception {
        Signature signature = Signature.getInstance("SHA256withRSA");
        signature.initSign(privateKey);
        signature.update(plainText.getBytes());
        return signature.sign();
    }

    /**
     * 使用公钥验证签名是否正确
     * @param plainText 文本
     * @param signatureBytes 签名后的字节数组
     * @param publicKey 公钥
     * @return 签名是否正确
     * @throws Exception 发生异常
     */
    public static boolean verify(String plainText, byte[] signatureBytes, PublicKey publicKey) throws Exception {
        Signature signature = Signature.getInstance("SHA256withRSA");
        signature.initVerify(publicKey);
        signature.update(plainText.getBytes());
        return signature.verify(signatureBytes);
    }

    /**
     * 将字符串格式的公钥转换成{@link PublicKey}格式
     * @param publicKey 公钥
     * @return {@link PublicKey}格式公钥
     * @throws Exception 发生异常
     */
    public static PublicKey loadPublicKey(String publicKey) throws Exception {
        byte[] data = Base64.getDecoder().decode((publicKey.getBytes()));
        X509EncodedKeySpec spec = new X509EncodedKeySpec(data);
        KeyFactory fact = KeyFactory.getInstance("RSA");
        return fact.generatePublic(spec);
    }

    /**
     * 将字符串格式私钥转转成{@link PrivateKey}格式
     * @param privateKey 私钥
     * @return {@link PrivateKey}格式私钥
     * @throws Exception 发生异常
     */
    public static PrivateKey loadPrivateKey(String privateKey) throws Exception {
        byte[] clear = Base64.getDecoder().decode(privateKey.getBytes());
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(clear);
        KeyFactory fact = KeyFactory.getInstance("RSA");
        PrivateKey priKey = fact.generatePrivate(keySpec);
        Arrays.fill(clear, (byte) 0);
        return priKey;
    }

    /**
     * 将公钥转换成字符串格式
     * @param publicKey 公钥
     * @return 转换后的字符串
     * @throws Exception 发生异常
     */
    public static String parsePublicKey(PublicKey publicKey) throws Exception {
        return Base64.getEncoder().encodeToString(publicKey.getEncoded());
    }

    /**
     * 将私钥转换成字符串格式
     * @param privateKey  私钥
     * @return 转换后的字符串
     * @throws Exception 发生异常
     */
    public static String parsePrivateKey(PrivateKey privateKey) throws Exception {
        return Base64.getEncoder().encodeToString(privateKey.getEncoded());
    }
}
